from typing import Annotated, TypedDict
from langgraph.graph import add_messages, StateGraph, START, END

from src.agent.my_llm import llm


class ChatState(TypedDict):
    messages: Annotated[list, add_messages]


def chat_bot(state: ChatState):
    return {"messages": [llm.invoke(state["messages"])]}


builder = StateGraph(ChatState)

builder.add_node(chat_bot, "chat_bot")
builder.add_edge(START, "chat_bot")
builder.add_edge("chat_bot", END)
graph = builder.compile()
